import pandas as pd
import os,glob,math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
from tqdm import tqdm
from global_land_mask import globe
from gsj_typhoon import tydat,see,count_rapidgrow,tydat_NH,average_datetime,split_str_id
from geopy.distance import geodesic
import matplotlib.ticker as ticker
import seaborn as sns

names,track_ids = ['2023NH42_36','2023NH47_46'],[42,47]
tynames,tyids = split_str_id(names)
draw_obs_opt = True

for trackid,ty,tyid in zip(track_ids,tynames,tyids):
    ''' init '''
    RIstd = 7
    dt = timedelta(days=1)
    directory =  f'D:\\met_data\\ty_ensemble\\{ty}_{tyid}\\'    # 要筛选的文件开头
    dates_name = os.listdir(directory)
    dates_paths = [directory + date for date in dates_name]
    obs_path = r'D:\met_data\ty_obs\storms\storms_NH2023'
    pic_savepath = rf'D:\met_data\pics\{ty}_{tyid}-RI{RIstd}\heatmap'
    os.makedirs(pic_savepath,exist_ok=True)
    
    
    ''' calc 2d numbers '''
    # 遍历所有起报时间
    max_delta_time = timedelta(days=0)
    for j,i in enumerate(dates_paths):
        mem_paths = glob.glob( os.path.join(i,"TRACK_ID_*") )
        # sorted_mem_paths = sorted(mem_paths, key=lambda x: int(x.split('_')[-1]))
        start_time = datetime.strptime(dates_name[j],"%Y%m%d%H")
        # 遍历所有成员
        for single_mem_path in mem_paths:
            try:
                single_mem = tydat(single_mem_path,RIstd) 
                delta_time = single_mem.time[-1] - start_time
                if delta_time>max_delta_time:
                    max_delta_time = delta_time
            except Exception as e:
                pass
    max_dt_num = math.ceil(max_delta_time/dt) 
    
    ''' 求各起报时间的图 '''
    tyobs = tydat_NH(obs_path,trackid)
    rcd_obs = count_rapidgrow(RIstd, tyobs.umax, tyobs.time)
    num_result = np.zeros(( max_dt_num,len(dates_name) ))
    dis_result = np.zeros(( max_dt_num,len(dates_name) ))
    dt_result  = np.zeros(( max_dt_num,len(dates_name) ))
    
    
    for i in tqdm( range(len(dates_paths)), desc=f'processing  {ty}'):
        '''读取这个起报时间的所有集合'''
        mem_paths = glob.glob( os.path.join(dates_paths[i],"TRACK_ID_*") )
        sorted_mem_paths = sorted(mem_paths, key=lambda x: int(x.split('_')[-1]))
        ty_mem = [tydat(i,RIstd) for i in sorted_mem_paths]
        start_dt = datetime.strptime( dates_name[i],"%Y%m%d%H" )
        for j in range(max_dt_num):
            dt = 1  # 指间隔1d求数据
            t0 = dt*j
            t1 = dt*j+dt
            t0_dt = start_dt + timedelta(days=t0)
            t1_dt = start_dt + timedelta(days=t1)
            '''如果观测在这两天没有RI，就直接跳过'''
            idx_obs_2d = (tyobs.time>=t0_dt)&(tyobs.time<t1_dt)
            idx_obs_2d_RI = (rcd_obs[idx_obs_2d]==1)
            if np.sum(rcd_obs[idx_obs_2d]) ==0.:
                ''' only record numbers here '''
                for k in range(len(ty_mem)):
                    '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''
                    idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)
                    idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)
                    num_result[j,i] += np.sum(idx_mem_2d_RI)
                continue
            else:
                '''如果2d内有多个RI,取平均'''
                time_obs_ave = average_datetime(tyobs.time[idx_obs_2d][idx_obs_2d_RI])
                lon_obs_ave  = np.mean(tyobs.lon[idx_obs_2d][idx_obs_2d_RI])
                lat_obs_ave  = np.mean(tyobs.lat[idx_obs_2d][idx_obs_2d_RI])
                '''开始计算'''
                num = 0 
                dis_sum = 0 
                dt_sum  = 0 
                for k in range(len(ty_mem)):
                    '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''
                    idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)
                    idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)
                    time_mem = ty_mem[k].time[idx_mem_2d][idx_mem_2d_RI]
                    lat_mem  = ty_mem[k].lat[idx_mem_2d][idx_mem_2d_RI]
                    lon_mem  = ty_mem[k].lon[idx_mem_2d][idx_mem_2d_RI]
                    if np.sum(idx_mem_2d_RI) ==0.:
                        continue
                    else:
                        num+=np.sum(idx_mem_2d_RI)
                        for l in range(len(lon_mem)):
                            dis_sum+=geodesic((lat_obs_ave,lon_obs_ave),(lat_mem[l],lon_mem[l])).kilometers
                            dt_sum+=abs((time_mem[l]-time_obs_ave).total_seconds())/3600
            if num==0:
                continue
            num_result[j,i] = num
            dis_result[j,i] = dis_sum/num
            dt_result[j,i]  = dt_sum/num #hour
            
    if draw_obs_opt == True : 
       ''' init obsCMA; RI points reord '''
       obs = tydat_NH(obs_path,trackid) 
       record_obs = count_rapidgrow(RIstd,obs.umax, obs.time)
       
       ''' prepare day array 1d: first start time of members '''
       draw_data_obs = np.zeros( (len(dates_name),max_dt_num),dtype=int )
       tt,dd = draw_data_obs.shape
       
       for t in range(tt):
           #预报时效
           for d in range(dd):
               try:
                   start_time = datetime.strptime(dates_name[t],"%Y%m%d%H")
                   cal_time = start_time + d*timedelta(days=1)
                   calend_time = cal_time + timedelta(days=1)
                   # 对齐
                   trange = (obs.time>=cal_time) & (obs.time<calend_time)
                   draw_data_obs[t][d] = 1 if np.sum(record_obs[trange])>0 else 0 
               except Exception as e:
                   pass           

    
    # see(num_result[::-1,:])
    # see(dis_result)
    # see(dt_result)
    
    var = dt_result
    plt.figure(figsize=(15, 6),dpi=1500)
    ''' heatmap '''
    ax = sns.heatmap(var[:15],annot=True,
                     annot_kws={'size': 4},fmt=".1f", cmap="YlGnBu", 
                     cbar_kws={'label': 'average RI time gap between obs and ensember members Units: Hour'}, 
                     linewidths=0.5)  #,vmin=0,vmax =30
    ax.set_ylabel("Forecast realative to start time", fontsize=12)
    ax.set_xlabel("Start time", fontsize=12)
    #  手动设置xticks数量,snsheatmap default上限是19个
    ax.set_xticks(np.arange(var.shape[1])) 
    ax.set_xticklabels(dates_name, rotation=45) 
    plt.yticks(rotation=0)
    ylabels=[str(i)+"d"  for i in range(1,16)]
    ax.set_yticklabels(ylabels)
    ax.set_title(ty+f"-RI{RIstd} average dt between RIobs and RImems" )
    cbar = ax.collections[0].colorbar
    cbar.locator = ticker.MaxNLocator(integer=True)
    cbar.update_ticks()
    
    
    ''' plot '''
    points = np.where(draw_data_obs[:,:15].T == 1)
    y_coords = points[0]+0.5
    x_coords = points[1]+0.15
    ax.scatter(x_coords, y_coords, color='red', marker='o',s=10)
    
    plt.savefig(f"{pic_savepath}\\heatmap_dt_days_startime",dpi=900)
    # plt.show()
    plt.close()
